package main

import (
	"encoding/base64"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os/exec"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/pion/webrtc/v3"
	"github.com/pion/webrtc/v3/pkg/media"
)

type Signal struct {
	SDP       string `json:"sdp"`
	Type      string `json:"type"`
	Candidate string `json:"candidate,omitempty"`
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Println("Server starting...")
}

func main() {
	http.HandleFunc("/ws", handleWebSocket)
	log.Println("Starting HTTP server on :8080")
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatalf("HTTP server failed: %v", err)
	}
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	log.Println("New WebSocket connection request from", r.RemoteAddr)
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade failed: %v", err)
		return
	}
	defer func() {
		log.Println("Closing WebSocket connection")
		conn.Close()
	}()

	// 初始化 WebRTC 配置
	log.Println("Initializing WebRTC configuration with STUN")
	config := webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{
			{URLs: []string{"stun:stun.l.google.com:19302"}},
			// 可选：添加 TURN 服务器
			// {URLs: []string{"turn:your.turn.server:3478"}, Username: "username", Credential: "password"},
		},
	}

	// 创建媒体引擎并注册 H.264 编解码器
	log.Println("Creating media engine and registering H.264 codec")
	mediaEngine := &webrtc.MediaEngine{}
	if err := mediaEngine.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{
			MimeType:    webrtc.MimeTypeH264,
			ClockRate:   90000,
			SDPFmtpLine: "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=42e01f",
		},
		PayloadType: 96,
	}, webrtc.RTPCodecTypeVideo); err != nil {
		log.Fatalf("Failed to register H264 codec: %v", err)
	}

	// 创建 API 和 PeerConnection
	log.Println("Creating WebRTC API and PeerConnection")
	api := webrtc.NewAPI(webrtc.WithMediaEngine(mediaEngine))
	peerConnection, err := api.NewPeerConnection(config)
	if err != nil {
		log.Fatalf("Failed to create peer connection: %v", err)
	}
	defer func() {
		log.Println("Closing PeerConnection")
		peerConnection.Close()
	}()

	// 创建视频轨道
	log.Println("Creating video track")
	videoTrack, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{
		MimeType:  webrtc.MimeTypeH264,
		ClockRate: 90000,
	}, "video", "pion")
	if err != nil {
		log.Fatalf("Failed to create video track: %v", err)
	}

	// 添加轨道到 PeerConnection
	log.Println("Adding video track to PeerConnection")
	rtpSender, err := peerConnection.AddTrack(videoTrack)
	if err != nil {
		log.Fatalf("Failed to add video track: %v", err)
	}

	// 处理 RTCP 反馈
	go func() {
		log.Println("Starting RTCP feedback loop")
		rtcpBuf := make([]byte, 1500)
		for {
			n, _, err := rtpSender.Read(rtcpBuf)
			if err != nil {
				log.Printf("RTCP read error: %v", err)
				return
			}
			log.Printf("Received %d bytes of RTCP data", n)
		}
	}()

	// 缓存 ICE 候选者
	pendingCandidates := make([]webrtc.ICECandidateInit, 0)
	var remoteDescSet bool

	// 处理 ICE 候选者
	peerConnection.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c != nil {
			log.Printf("New ICE candidate: %s", c.ToJSON().Candidate)
			candidate, _ := json.Marshal(c.ToJSON())
			if err := conn.WriteJSON(Signal{
				Candidate: string(candidate),
				Type:      "candidate",
			}); err != nil {
				log.Printf("Failed to send ICE candidate: %v", err)
			}
		}
	})

	// 处理 ICE 连接状态变化
	peerConnection.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
		log.Printf("ICE connection state changed: %s", state.String())
		if state == webrtc.ICEConnectionStateFailed {
			log.Println("ICE connection failed, closing PeerConnection")
			peerConnection.Close()
			conn.WriteJSON(Signal{Type: "error", SDP: "ICE connection failed"})
		}
	})

	// 处理 WebSocket 消息
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		log.Println("Starting WebSocket message loop")
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				log.Printf("WebSocket read error: %v", err)
				return
			}
			log.Printf("Received WebSocket message: %s", string(msg))

			var signal Signal
			if err := json.Unmarshal(msg, &signal); err != nil {
				log.Printf("JSON unmarshal error: %v", err)
				continue
			}

			switch signal.Type {
			case "offer":
				log.Println("Received offer from client")
				sdp, err := base64.StdEncoding.DecodeString(signal.SDP)
				if err != nil {
					log.Printf("SDP decode error: %v", err)
					continue
				}
				log.Printf("Decoded Offer SDP: %s", string(sdp)[:200]+"...")

				if err := peerConnection.SetRemoteDescription(webrtc.SessionDescription{
					Type: webrtc.SDPTypeOffer,
					SDP:  string(sdp),
				}); err != nil {
					log.Printf("Set remote description error: %v", err)
					continue
				}
				log.Println("Set remote description successfully")
				remoteDescSet = true

				// 处理缓存的 ICE 候选者
				for _, candidate := range pendingCandidates {
					if err := peerConnection.AddICECandidate(candidate); err != nil {
						log.Printf("Add cached ICE candidate error: %v", err)
					} else {
						log.Println("Added cached ICE candidate successfully")
					}
				}
				pendingCandidates = nil

				log.Println("Creating answer")
				answer, err := peerConnection.CreateAnswer(nil)
				if err != nil {
					log.Printf("Create answer error: %v", err)
					continue
				}

				log.Println("Setting local description")
				if err := peerConnection.SetLocalDescription(answer); err != nil {
					log.Printf("Set local description error: %v", err)
					continue
				}

				log.Printf("Answer SDP: %s", answer.SDP[:200]+"...")
				answerSDP := base64.StdEncoding.EncodeToString([]byte(answer.SDP))
				if err := conn.WriteJSON(Signal{
					SDP:  answerSDP,
					Type: "answer",
				}); err != nil {
					log.Printf("WebSocket write error: %v", err)
					continue
				}
				log.Println("Sent answer to client")

				log.Println("Waiting for ICE gathering to complete")
				gatherComplete := webrtc.GatheringCompletePromise(peerConnection)
				<-gatherComplete
				log.Println("ICE gathering completed")

				log.Println("Starting FFmpeg stream")
				go startFFmpeg(videoTrack)

			case "candidate":
				log.Println("Received ICE candidate from client")
				var candidate webrtc.ICECandidateInit
				if err := json.Unmarshal([]byte(signal.Candidate), &candidate); err != nil {
					log.Printf("ICE candidate unmarshal error: %v", err)
					continue
				}
				if !remoteDescSet {
					log.Println("Caching ICE candidate, remote description not set")
					pendingCandidates = append(pendingCandidates, candidate)
				} else {
					if err := peerConnection.AddICECandidate(candidate); err != nil {
						log.Printf("Add ICE candidate error: %v", err)
						continue
					}
					log.Println("Added ICE candidate successfully")
				}
			default:
				log.Printf("Unknown message type: %s", signal.Type)
			}
		}
	}()

	wg.Wait()
}

func startFFmpeg(track *webrtc.TrackLocalStaticSample) {
	log.Println("Preparing FFmpeg command")
	ffmpegCmd := exec.Command("ffmpeg",
		"-re", // 实时读取
		"-i", "/home/weiqiangren/Documents/big-h264.mp4",
		"-c:v", "libx264", // 重新编码以控制参数
		"-r", "25", // 强制 25fps
		"-crf", "23", // 恒定质量，23 为高质量
		"-b:v", "400k", // 目标比特率 400kbps
		"-maxrate", "450k", // 最大比特率
		"-bufsize", "800k", // 缓冲区大小
		"-g", "15", // 每 0.6 秒一个关键帧（25fps * 0.6）
		"-keyint_min", "15", // 最小关键帧间隔
		"-sc_threshold", "0", // 禁用场景检测
		"-preset", "fast", // 提高编码质量
		"-tune", "zerolatency", // 零延迟优化
		"-pix_fmt", "yuv420p", // 强制 YUV 4:2:0 色空间
		"-bsf:v", "h264_mp4toannexb",
		"-f", "h264",
		"pipe:1",
	)

	// 启用详细日志
	ffmpegCmd.Args = append(ffmpegCmd.Args, "-loglevel", "verbose")

	stderr, err := ffmpegCmd.StderrPipe()
	if err != nil {
		log.Printf("FFmpeg stderr pipe error: %v", err)
		return
	}
	go func() {
		buf := make([]byte, 1024)
		for {
			n, err := stderr.Read(buf)
			if err != nil {
				if err != io.EOF {
					log.Printf("FFmpeg stderr read error: %v", err)
				}
				return
			}
			log.Printf("FFmpeg stderr: %s", string(buf[:n]))
		}
	}()

	log.Println("Creating FFmpeg stdout pipe")
	stdout, err := ffmpegCmd.StdoutPipe()
	if err != nil {
		log.Fatalf("FFmpeg stdout pipe error: %v", err)
	}

	log.Println("Starting FFmpeg process")
	if err := ffmpegCmd.Start(); err != nil {
		log.Fatalf("FFmpeg start error: %v", err)
	}

	log.Println("Starting FFmpeg data streaming")
	buf := make([]byte, 1024*1024)
	for {
		n, err := stdout.Read(buf)
		if err != nil {
			if err == io.EOF {
				log.Println("FFmpeg stream ended (EOF)")
				break
			}
			log.Printf("FFmpeg read error: %v", err)
			continue
		}
		log.Printf("Read %d bytes from FFmpeg", n)

		if err := track.WriteSample(media.Sample{
			Data:     buf[:n],
			Duration: time.Second / 25,
		}); err != nil {
			log.Printf("Track write error: %v", err)
			break
		}
		log.Printf("Wrote %d bytes to WebRTC track", n)
	}

	log.Println("Waiting for FFmpeg process to exit")
	if err := ffmpegCmd.Wait(); err != nil {
		log.Printf("FFmpeg wait error: %v", err)
	}
	log.Println("FFmpeg process exited")
}
